import os
import sys
import json
from flask import Flask, render_template, request, Response, stream_with_context, flash, redirect, url_for
import requests
from werkzeug.datastructures import FileStorage

# 导入 cozepy 相关的模块
# 请确保已安装: pip install cozepy
try:
    from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType
except ImportError:
    print("错误: 'cozepy' 库未安装。请运行 'pip install cozepy' 进行安装。", file=sys.stderr)
    sys.exit(1)

# --- 1. Flask 应用与配置 ---
app = Flask(__name__)
# 强烈建议在生产环境中使用环境变量或专门的配置文件
app.secret_key = os.getenv("FLASK_SECRET_KEY", "a_strong_dev_secret_key")

# --- Coze API 配置 ---
# 优先从环境变量获取，如果未设置则使用下面的默认值
# 强烈建议将您的令牌设置为环境变量，以避免暴露在代码中
COZE_API_TOKEN = os.getenv("COZE_API_TOKEN", "上传密钥")
token="工作流密钥"
WORKFLOW_ID = os.getenv("COZE_WORKFLOW_ID", "工作流ID")
API_BASE = os.getenv("COZE_API_BASE", COZE_CN_BASE_URL)
UPLOAD_URL = f"{API_BASE}/v1/files/upload"


# --- 2. 核心功能函数 ---

def upload_file_data_to_coze(filename: str, file_content: bytes, mimetype: str):

    if not COZE_API_TOKEN or "pat_" not in COZE_API_TOKEN:
        raise ValueError("Coze API Token 配置无效或缺失。")

    headers = {"Authorization": f"Bearer {COZE_API_TOKEN}"}
    files_payload = {"file": (filename, file_content, mimetype)}

    response = requests.post(UPLOAD_URL, headers=headers, files=files_payload)
    response.raise_for_status()  # 如果HTTP状态码为4xx或5xx，则抛出异常

    response_data = response.json()
    if response_data.get("code") == 0 and "data" in response_data and "id" in response_data["data"]:
        return response_data["data"]["id"]
    else:
        error_msg = response_data.get("msg", "未知上传错误")
        raise Exception(f"Coze API 上传文件 '{filename}' 失败: {error_msg}")


def handle_workflow_stream_generator(client: Coze, workflow_id: str, stream: Stream[WorkflowEvent]):
    for event in stream:
        if event.event == WorkflowEventType.MESSAGE:
            yield f"💬 [消息]: {event.message.content}\n"
        elif event.event == WorkflowEventType.ERROR:
            yield f"❌ [错误]: {event.error}\n"
        elif event.event == WorkflowEventType.INTERRUPT:
            yield f"⏸️ [中断]: 工作流需要额外输入。将自动回复 'hey' 以继续...\n"
            # 自动回复以继续流程
            resume_data = "hey"
            try:
                new_stream = client.workflows.runs.resume(
                    workflow_id=workflow_id,
                    event_id=event.interrupt.interrupt_data.event_id,
                    resume_data=resume_data,
                    interrupt_type=event.interrupt.interrupt_data.type,
                )
                # 递归地从新的事件流中 yield
                yield from handle_workflow_stream_generator(client, workflow_id, new_stream)
            except Exception as e:
                yield f"❌ [恢复中断失败]: {e}\n"
        elif event.event == WorkflowEventType.DONE:
            yield "\n🎉 [完成]: 工作流执行完毕。\n"
            break


# --- 3. Flask 路由 ---

@app.route('/', methods=['GET', 'POST'])
def index():

    if request.method == 'POST':
        files = request.files.getlist('files')

        # --- 表单验证 ---
        if not files or not files[0].filename:
            flash('请至少选择一个文件！', 'error')
            return redirect(url_for('index'))
        if not WORKFLOW_ID:
            flash('错误：服务器未配置工作流ID (COZE_WORKFLOW_ID)。', 'error')
            return redirect(url_for('index'))
        if not COZE_API_TOKEN or "pat_" not in COZE_API_TOKEN:
            flash('错误：服务器未配置有效的Coze API令牌 (COZE_API_TOKEN)。', 'error')
            return redirect(url_for('index'))

        # --- 关键重构：在启动流式响应前，预先读取所有文件 ---
        # 这种方法可以彻底避免在生成器内部处理文件流时可能遇到的上下文或流关闭问题。
        try:
            files_to_upload = []
            for file in files:
                files_to_upload.append({
                    "filename": file.filename,
                    "content": file.read(),  # 立即将文件内容读入内存
                    "mimetype": file.mimetype
                })
        except Exception as e:
            flash(f"读取上传的文件时出错: {e}", 'error')
            return redirect(url_for('index'))

        def generate_response_stream():
            """
            用于流式响应的生成器。
            它会执行文件上传和工作流调用，并逐步 yield 结果。
            """
            file_ids = []
            try:
                # --- 步骤 1: 上传文件 ---
                yield "开始处理请求...\n"
                for file_data in files_to_upload:  # 遍历已读入内存的文件数据
                    yield f"  - 正在上传文件: {file_data['filename']} ...\n"
                    # 对每个文件调用新的上传函数
                    file_id = upload_file_data_to_coze(
                        filename=file_data['filename'],
                        file_content=file_data['content'],
                        mimetype=file_data['mimetype']
                    )
                    file_ids.append(file_id)
                    yield f"  - ✅ 文件上传成功！File ID: {file_id}\n"

                yield "\n所有文件上传完毕。\n"

                # --- 步骤 2: 准备并调用工作流 ---
                yield "🚀 准备调用工作流...\n"

                input_params = [json.dumps({"file_id": fid}) for fid in file_ids]

                parameters = {
                    "input": input_params,
                    "prompt": ""
                }

                client = Coze(auth=TokenAuth(token=token), base_url=API_BASE)

                workflow_stream = client.workflows.runs.stream(
                    workflow_id=WORKFLOW_ID,
                    parameters=parameters,
                )
                yield "工作流已启动，等待事件流...\n\n"

                # --- 步骤 3: 处理事件流 ---
                yield from handle_workflow_stream_generator(client, WORKFLOW_ID, workflow_stream)

            except Exception as e:
                yield f"\n❌ [严重错误]: 操作失败: {e}\n"

        # 返回一个流式响应
        return Response(stream_with_context(generate_response_stream()), mimetype='text/plain; charset=utf-8')

    # 对于GET请求，只渲染主页面
    return render_template('index.html')


if __name__ == '__main__':
    # 建议在生产环境中使用 Gunicorn 或 uWSGI 等WSGI服务器
    app.run(debug=True, port=5001)

